package com.zys.background.common.config;

import feign.RequestInterceptor;
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
import org.springframework.boot.web.servlet.ServletRegistrationBean;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;
import org.springframework.web.servlet.DispatcherServlet;

import javax.servlet.http.HttpServletRequest;
import java.util.Enumeration;

/**
 * @apiNote 给 Feign 添加上请求头
 * @author zys
 * @since 2022-12-01
 */
@Configuration
@EnableAutoConfiguration
public class FeignConfiguration {
  /**
   * 让DispatcherServlet向子线程传递RequestContext
   *
   * @param servlet servlet
   * @return 注册bean
   */
  @Bean
  public ServletRegistrationBean<DispatcherServlet> dispatcherRegistration(DispatcherServlet servlet) {
    servlet.setThreadContextInheritable(true);
    return new ServletRegistrationBean<>(servlet, "/**");
  }

  /**
   * 覆写拦截器, 在 feign 发送请求前取出原来的 header 并转发
   *
   * @return 拦截器
   */
  @Bean
  public RequestInterceptor requestInterceptor() {
    return template -> {
      RequestAttributes requestAttributes = RequestContextHolder.getRequestAttributes();
      if (requestAttributes != null) {
        ServletRequestAttributes attributes = (ServletRequestAttributes) requestAttributes;
        HttpServletRequest request = attributes.getRequest();
        // 获取请求头
        Enumeration<String> headerNames = request.getHeaderNames();
        if (headerNames != null) {
          while (headerNames.hasMoreElements()) {
            String name = headerNames.nextElement();
            String values = request.getHeader(name);
            // 将请求头保存到模板中
            template.header(name, values);
          }
        }
      }
    };
  }
}
